import torch
from torchvision import models, transforms
from PIL import Image
import torch.nn as nn

# 设置计算设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 标签映射：模型输出标签 → 中文疾病名称
label_map = {
    'N': '正常',
    'D': '糖尿病',
    'G': '青光眼',
    'C': '白内障',
    'A': 'AMD',
    'H': '高血压',
    'M': '近视',
    'O': '其他异常'
}

# 类别标签列表（需与模型训练时顺序一致）
classes = list(label_map.keys())

# 图像预处理流程
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 统一尺寸
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

# 加载 ResNet50 模型结构
model = models.resnet50(weights=None)
model.fc = nn.Sequential(
    nn.Linear(model.fc.in_features, len(classes)),
    nn.Sigmoid()
)

# 加载模型参数
model.load_state_dict(torch.load("resnet50_8class_multilabel.pth", map_location=device))
model = model.to(device)
model.eval()

def predict_image(image_path):
    """给定图片路径，返回预测类别与概率"""
    img = Image.open(image_path).convert("RGB")
    img_tensor = transform(img).unsqueeze(0).to(device)

    with torch.no_grad():
        output = model(img_tensor)
        pred = output.argmax(dim=1).item()
        probs = torch.nn.functional.softmax(output, dim=1).cpu().numpy()[0]

    result = {
        "prediction": label_map[classes[pred]],
        "probs": {label_map[cls]: round(probs[i]*100, 2) for i, cls in enumerate(classes)}
    }
    return result



# # 英文标签
# classes = ['A', 'C', 'D', 'G', 'H', 'M', 'N', 'O']
# # 对应中文标签
# class_map = {
#     'N': '正常',
#     'D': '糖尿病',
#     'G': '青光眼',
#     'C': '白内障',
#     'A': 'AMD',
#     'H': '高血压',
#     'M': '近视',
#     'O': '其他异常'
# }
